import os
import numpy as np
import torch
from tensorboardX import SummaryWriter
from utils.shared_buffer import SharedReplayBuffer

def _t2n(x):
    """将PyTorch张量转换为NumPy数组。"""
    return x.detach().cpu().numpy()

class Runner(object):
    """
    训练循环策略的基类。
    :param config: (dict) 包含训练参数的配置字典。
    """
    def __init__(self, config):
        """
        初始化Runner类
        
        参数:
            config: 配置字典，包含训练所需的所有参数和环境
        """
        self.all_args = config['all_args']  # 所有参数
        self.envs = config['envs']  # 训练环境
        self.eval_envs = config['eval_envs']  # 评估环境
        self.device = config['device']  # 计算设备
        self.num_agents = config['num_agents']  # 智能体数量
        if config.__contains__("render_envs"):
            self.render_envs = config['render_envs']  # 渲染环境

        # 参数
        self.env_name = self.all_args.env_name  # 环境名称
        self.algorithm_name = self.all_args.algorithm_name  # 算法名称
        self.experiment_name = self.all_args.experiment_name  # 实验名称
        self.use_centralized_V = self.all_args.use_centralized_V  # 是否使用中心化价值函数
        self.use_obs_instead_of_state = self.all_args.use_obs_instead_of_state  # 是否使用观察代替状态
        self.num_env_steps = self.all_args.num_env_steps  # 环境步数
        self.episode_length = self.all_args.episode_length  # 回合长度
        self.n_rollout_threads = self.all_args.n_rollout_threads  # 并行环境数量
        self.n_eval_rollout_threads = self.all_args.n_eval_rollout_threads  # 评估并行环境数量
        self.n_render_rollout_threads = self.all_args.n_render_rollout_threads  # 渲染并行环境数量
        self.use_linear_lr_decay = self.all_args.use_linear_lr_decay  # 是否使用线性学习率衰减
        self.hidden_size = self.all_args.hidden_size  # 隐藏层大小
        self.use_render = self.all_args.use_render  # 是否渲染
        self.recurrent_N = self.all_args.recurrent_N  # 循环层数

        # 间隔
        self.save_interval = self.all_args.save_interval  # 保存间隔
        self.use_eval = self.all_args.use_eval  # 是否使用评估
        self.eval_interval = self.all_args.eval_interval  # 评估间隔
        self.log_interval = self.all_args.log_interval  # 日志间隔

        # 目录
        self.model_dir = self.all_args.model_dir  # 模型目录

        self.run_dir = config["run_dir"]  # 运行目录
        self.log_dir = str(self.run_dir / 'logs')  # 日志目录
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)  # 创建日志目录
        self.writter = SummaryWriter(self.log_dir)  # 创建TensorBoard写入器
        self.save_dir = str(self.run_dir / 'models')  # 保存目录
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)  # 创建保存目录

        from algorithms.algorithm.r_mappo import RMAPPO as TrainAlgo
        from algorithms.algorithm.rMAPPOPolicy import RMAPPOPolicy as Policy

        # 如果使用中心化价值函数，则使用共享观察空间，否则使用观察空间
        share_observation_space = self.envs.share_observation_space[0] if self.use_centralized_V else self.envs.observation_space[0]

        # 策略网络
        self.policy = Policy(self.all_args,
                            self.envs.observation_space[0],
                            share_observation_space,
                            self.envs.action_space[0],
                            device = self.device)

        if self.model_dir is not None:
            self.restore()  # 恢复模型

        # 算法
        self.trainer = TrainAlgo(self.all_args, self.policy, device = self.device)
        
        # 缓冲区
        self.buffer = SharedReplayBuffer(self.all_args,
                                        self.num_agents,
                                        self.envs.observation_space[0],
                                        share_observation_space,
                                        self.envs.action_space[0])

    def run(self):
        """收集训练数据，执行训练更新，并评估策略。"""
        raise NotImplementedError

    def warmup(self):
        """收集预训练的热身数据。"""
        raise NotImplementedError

    def collect(self, step):
        """收集用于训练的数据。"""
        raise NotImplementedError

    def insert(self, data):
        """
        将数据插入缓冲区。
        :param data: (Tuple) 要插入训练缓冲区的数据。
        """
        raise NotImplementedError
    
    @torch.no_grad()
    def compute(self):
        """计算收集数据的回报。"""
        self.trainer.prep_rollout()  # 准备收集数据
        next_values = self.trainer.policy.get_values(np.concatenate(self.buffer.share_obs[-1]),
                                                np.concatenate(self.buffer.rnn_states_critic[-1]),
                                                np.concatenate(self.buffer.masks[-1]))  # 获取下一个状态的价值
        next_values = np.array(np.split(_t2n(next_values), self.n_rollout_threads))  # 分割价值
        self.buffer.compute_returns(next_values, self.trainer.value_normalizer)  # 计算回报
    
    def train(self):
        """使用缓冲区中的数据训练策略。"""
        self.trainer.prep_training()  # 准备训练
        train_infos = self.trainer.train(self.buffer)  # 训练策略
        self.buffer.after_update()  # 更新后处理
        return train_infos  # 返回训练信息

    def save(self):
        """保存策略的演员和评论家网络。"""
        policy_actor = self.trainer.policy.actor  # 获取策略的演员网络
        torch.save(policy_actor.state_dict(), str(self.save_dir) + "/actor.pt")  # 保存演员网络
        policy_critic = self.trainer.policy.critic  # 获取策略的评论家网络
        torch.save(policy_critic.state_dict(), str(self.save_dir) + "/critic.pt")  # 保存评论家网络

    def restore(self):
        """从保存的模型中恢复策略的网络。"""
        policy_actor_state_dict = torch.load(str(self.model_dir) + '/actor.pt')  # 加载演员网络状态字典
        self.policy.actor.load_state_dict(policy_actor_state_dict)  # 加载演员网络
        if not self.all_args.use_render:
            policy_critic_state_dict = torch.load(str(self.model_dir) + '/critic.pt')  # 加载评论家网络状态字典
            self.policy.critic.load_state_dict(policy_critic_state_dict)  # 加载评论家网络
 
    def log_train(self, train_infos, total_num_steps):
        """
        记录训练信息。
        :param train_infos: (dict) 关于训练更新的信息。
        :param total_num_steps: (int) 训练环境步骤的总数。
        """
        for k, v in train_infos.items():
            self.writter.add_scalars(k, {k: v}, total_num_steps)  # 添加标量

    def log_env(self, env_infos, total_num_steps):
        """
        记录环境信息。
        :param env_infos: (dict) 关于环境状态的信息。
        :param total_num_steps: (int) 训练环境步骤的总数。
        """
        for k, v in env_infos.items():
            if len(v)>0:
                self.writter.add_scalars(k, {k: np.mean(v)}, total_num_steps)  # 添加标量
